import numpy as np
from typing import Dict, List, Tuple
import os
import logging
from datetime import datetime
from physics import PhysicsModel, PhaseType
from parallel_compute import (
    calculate_mass_flux_array,
    calculate_momentum_flux_array,
    calculate_wall_friction_array,
    calculate_interfacial_friction_array,
    calculate_lateral_flow_resistance_array,
    calculate_turbulent_mixing_array,
    solve_momentum_equation_parallel,
    solve_continuity_equation_parallel
)

class Solver:
    """两相流求解器类（使用并行计算）"""
    def __init__(self, grid, physics_model):
        self.grid = grid
        self.physics = physics_model
        
        # 求解器参数
        self.max_iterations = 1000
        self.convergence_tolerance = 1e-6
        self.relaxation_factors = {
            'pressure': 0.3,
            'velocity': 0.7,
            'void_fraction': 0.3
        }
        
        # CFL条件相关
        self.max_cfl = 0.8
        self.min_dt = 1e-6
        self.max_dt = 0.01
        
        # 时间步长
        self.dt = self.min_dt
        
        # 设置日志记录
        self._setup_logging()
        
        # 设置初始条件
        self.initialize_fields()
        
        # 添加PKM求解器
        self.pkm_solver = PKMSolver(grid, physics_model)
        
    def _setup_logging(self):
        """设置日志记录"""
        # 创建results文件夹（如果不存在）
        results_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'results')
        if not os.path.exists(results_dir):
            os.makedirs(results_dir)
        
        # 创建日志文件名（使用时间戳）
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        log_file = os.path.join(results_dir, f'simulation_{timestamp}.log')
        
        # 配置日志记录器
        self.logger = logging.getLogger('TwoFluidSCA')
        self.logger.setLevel(logging.INFO)
        
        # 文件处理器
        file_handler = logging.FileHandler(log_file, encoding='utf-8')
        file_handler.setLevel(logging.INFO)
        
        # 控制台处理器
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        
        # 创建格式化器
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        file_handler.setFormatter(formatter)
        console_handler.setFormatter(formatter)
        
        # 添加处理器
        self.logger.addHandler(file_handler)
        self.logger.addHandler(console_handler)
        
        self.logger.info(f"日志文件创建于: {log_file}")
        
    def log_and_print(self, message: str, level: str = 'info'):
        """同时记录日志并打印到控制台
        
        Args:
            message: 要记录的消息
            level: 日志级别 ('info', 'warning', 'error', 'debug')
        """
        if level == 'info':
            self.logger.info(message)
        elif level == 'warning':
            self.logger.warning(message)
        elif level == 'error':
            self.logger.error(message)
        elif level == 'debug':
            self.logger.debug(message)
        
    def initialize_fields(self):
        """初始化所有场量"""
        self.log_and_print("\n初始化计算场...")
        for subchannel in self.grid.subchannels:
            n_axial = len(subchannel.alpha)
            
            # 1. 压力场初始化：线性分布（常压，略高于大气压）
            P_in = 1.2e5  # 入口压力 (1.2 bar)
            P_out = 1.0e5  # 出口压力 (1.0 bar)
            subchannel.pressure = np.linspace(P_in, P_out, n_axial)
            subchannel.previous_pressure = subchannel.pressure.copy()
            
            # 2. 空泡份额初始化：入口值延拓（减小初始空泡份额以提高稳定性）
            alpha_in = 0.1  # 入口空泡份额（降低初始值）
            subchannel.alpha = np.full(n_axial, alpha_in)
            subchannel.previous_alpha = subchannel.alpha.copy()
            
            # 3. 速度场初始化（降低初始速度以提高稳定性）
            v_l_in = 0.15  # 入口液相速度
            v_g_in = 0.20  # 入口气相速度（略快于液相）
            
            # 初始化速度场（包括轴向和横向分量）
            subchannel.velocity = {
                'liquid': {
                    'axial': np.full(n_axial, v_l_in),
                    'lateral': np.zeros(n_axial)
                },
                'vapor': {
                    'axial': np.full(n_axial, v_g_in),
                    'lateral': np.zeros(n_axial)
                }
            }
            
            # 初始化previous_velocity
            subchannel.previous_velocity = {
                'liquid': {
                    'axial': subchannel.velocity['liquid']['axial'].copy(),
                    'lateral': subchannel.velocity['liquid']['lateral'].copy()
                },
                'vapor': {
                    'axial': subchannel.velocity['vapor']['axial'].copy(),
                    'lateral': subchannel.velocity['vapor']['lateral'].copy()
                }
            }
            
            # 初始化质量通量
            subchannel.previous_mass_flux = {
                'vapor': np.zeros(n_axial),
                'liquid': np.zeros(n_axial)
            }
            
            # 初始化漂移通量
            subchannel.previous_drift_flux = np.zeros(n_axial)
            
        self.log_and_print("场初始化完成：")
        self.log_and_print(f"- 压力范围：{P_in/1e5:.2f} bar - {P_out/1e5:.2f} bar")
        self.log_and_print(f"- 入口空泡份额：{alpha_in:.3f}")
        self.log_and_print(f"- 入口速度：液相 {v_l_in:.2f} m/s, 气相 {v_g_in:.2f} m/s")
        
    def calculate_time_step(self) -> float:
        """计算自适应时间步长"""
        # 计算最大速度
        max_velocity = 1e-6  # 设置一个最小速度以避免除零
        for subchannel in self.grid.subchannels:
            for phase in ['liquid', 'vapor']:
                max_velocity = max(max_velocity,
                                 np.max(np.abs(subchannel.velocity[phase]['axial'])),
                                 np.max(np.abs(subchannel.velocity[phase]['lateral'])))
        
        # 基于CFL条件计算时间步长
        dt_cfl = self.max_cfl * min(self.grid.dx, self.grid.dz) / max_velocity
        
        # 限制时间步长在允许范围内
        self.dt = np.clip(dt_cfl, self.min_dt, self.max_dt)
        
        return self.dt
    
    def predict_velocity(self, subchannel, dt: float):
        """速度预测步骤"""
        n_axial = len(subchannel.alpha)
        
        for i in range(1, n_axial-1):
            # 获取物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 使用动量方程预测速度
            for phase, props in [('vapor', props_g), ('liquid', props_l)]:
                # 计算压力梯度
                dp_dz = (subchannel.pressure[i+1] - subchannel.pressure[i-1]) / (2 * self.grid.dz)
                
                # 预测速度
                subchannel.velocity[phase]['axial'][i] = (
                    subchannel.previous_velocity[phase]['axial'][i] -
                    dt * dp_dz / props['density']
                )
    
    def solve_pressure_correction(self, subchannel):
        """求解压力修正方程"""
        n_axial = len(subchannel.alpha)
        
        # 构建压力修正方程系数矩阵
        A = np.zeros((n_axial, n_axial))
        b = np.zeros(n_axial)
        
        # 设置入口边界条件（固定压力）
        A[0, 0] = 1.0
        b[0] = 0.0  # 入口压力不需要修正
        
        # 设置出口边界条件（固定压力）
        A[-1, -1] = 1.0
        b[-1] = 0.0  # 出口压力不需要修正
        
        # 构建内部节点的方程
        for i in range(1, n_axial-1):
            # 获取物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 计算密度项（添加保护）
            rho_g = max(props_g['density'], 1e-6)
            rho_l = max(props_l['density'], 1e-6)
            
            # 计算速度项
            v_g = subchannel.velocity['vapor']['axial'][i]
            v_l = subchannel.velocity['liquid']['axial'][i]
            
            # 计算系数（考虑密度和空泡份额的影响，添加保护）
            alpha = np.clip(subchannel.alpha[i], 1e-6, 1-1e-6)
            coeff = (alpha / rho_g + (1 - alpha) / rho_l)
            
            # 添加数值稳定项
            stability_factor = 1e-8
            
            # 设置系数矩阵
            A[i, i-1] = -coeff / (self.grid.dz ** 2)
            A[i, i] = 2 * coeff / (self.grid.dz ** 2) + stability_factor
            A[i, i+1] = -coeff / (self.grid.dz ** 2)
            
            # 计算质量不平衡作为源项
            G_g = alpha * rho_g * v_g
            G_l = (1 - alpha) * rho_l * v_l
            
            # 计算质量流量的空间导数（使用中心差分）
            alpha_prev = np.clip(subchannel.alpha[i-1], 1e-6, 1-1e-6)
            v_g_prev = subchannel.velocity['vapor']['axial'][i-1]
            v_l_prev = subchannel.velocity['liquid']['axial'][i-1]
            
            G_g_prev = alpha_prev * rho_g * v_g_prev
            G_l_prev = (1 - alpha_prev) * rho_l * v_l_prev
            
            dG_g_dz = (G_g - G_g_prev) / self.grid.dz
            dG_l_dz = (G_l - G_l_prev) / self.grid.dz
            
            b[i] = -(dG_g_dz + dG_l_dz)
        
        try:
            # 添加对角优势以提高稳定性
            for i in range(n_axial):
                A[i,i] = max(A[i,i], 1.1 * np.sum(np.abs(A[i,:])))
            
            # 使用直接求解器求解压力修正方程
            pressure_correction = np.linalg.solve(A, b)
            
            # 检查解的有效性
            if not np.all(np.isfinite(pressure_correction)):
                raise ValueError("压力修正解包含无效值")
            
            # 限制压力修正的幅度
            max_correction = 0.1 * np.mean(subchannel.pressure)
            pressure_correction = np.clip(pressure_correction, -max_correction, max_correction)
            
            # 应用压力修正（带松弛因子）
            subchannel.pressure += self.relaxation_factors['pressure'] * pressure_correction
            
            # 确保压力始终在物理范围内
            subchannel.pressure = np.maximum(subchannel.pressure, 1e5)  # 最小压力为1个大气压
            
        except (np.linalg.LinAlgError, ValueError) as e:
            print(f"警告：压力修正方程求解失败: {str(e)}")
            # 使用简单的压力梯度估计
            dp = (subchannel.pressure[-1] - subchannel.pressure[0]) / (n_axial - 1)
            subchannel.pressure = np.linspace(subchannel.pressure[0], subchannel.pressure[-1], n_axial)
    
    def correct_velocity(self, subchannel):
        """基于压力修正更新速度场"""
        n_axial = len(subchannel.alpha)
        
        for i in range(1, n_axial-1):
            for phase in ['vapor', 'liquid']:
                # 计算速度修正
                dp_correction = (subchannel.pressure[i+1] - subchannel.pressure[i-1]) / (2 * self.grid.dz)
                velocity_correction = -dp_correction * self.relaxation_factors['velocity']
                
                # 更新速度
                subchannel.velocity[phase]['axial'][i] += velocity_correction
    
    def solve_subchannel_parallel(self, subchannel, dt: float):
        """并行求解单个子通道，包含横向流动和空隙漂移"""
        n_axial = len(subchannel.alpha)
        
        # 获取相邻子通道信息
        adjacent_channels = self.grid.get_adjacent_channels(subchannel.id)
        
        # 获取几何参数
        geom = subchannel.get_geometry_parameters()
        
        # 准备数组数据
        alpha = subchannel.alpha
        pressure = subchannel.pressure
        v_g_axial = subchannel.velocity['vapor']['axial']
        v_l_axial = subchannel.velocity['liquid']['axial']
        v_g_lateral = subchannel.velocity['vapor']['lateral']
        v_l_lateral = subchannel.velocity['liquid']['lateral']
        
        # 获取物性
        props_g = self.physics.calculate_properties(pressure[0], PhaseType.VAPOR)
        props_l = self.physics.calculate_properties(pressure[0], PhaseType.LIQUID)
        rho_g = props_g['density']
        rho_l = props_l['density']
        
        # 计算空隙漂移速度
        v_gj, C0 = self.physics.calculate_drift_velocity(rho_g, rho_l)
        
        # 修正气相速度（考虑漂移速度）
        j = alpha * v_g_axial + (1 - alpha) * v_l_axial  # 混合体积通量
        v_g_axial = C0 * j + v_gj
        
        # 轴向动量方程
        for i in range(1, n_axial-1):
            # 计算轴向压力梯度
            dp_dz = (pressure[i+1] - pressure[i-1]) / (2 * self.grid.dz)
            
            # 计算轴向壁面摩擦力
            F_w_g = self.physics.calculate_wall_friction(
                PhaseType.VAPOR, rho_g, v_g_axial[i],
                geom['hydraulic_diameter'], props_g['viscosity']
            )
            F_w_l = self.physics.calculate_wall_friction(
                PhaseType.LIQUID, rho_l, v_l_axial[i],
                geom['hydraulic_diameter'], props_l['viscosity']
            )
            
            # 计算界面摩擦力
            F_i = self.physics.calculate_interfacial_friction(
                alpha[i], rho_g, rho_l, v_g_axial[i], v_l_axial[i],
                geom['hydraulic_diameter']
            )
            
            # 横向动量交换（遍历相邻通道）
            M_lat_g_total = 0
            M_lat_l_total = 0
            for adj_ch in adjacent_channels:
                # 获取相邻通道数据
                v_g_adj = adj_ch.velocity['vapor']['axial'][i]
                v_l_adj = adj_ch.velocity['liquid']['axial'][i]
                
                # 获取连接参数
                conn_params = subchannel.get_connection_parameters(adj_ch)
                
                # 计算横向动量交换
                M_lat_g, M_lat_l = self.physics.calculate_lateral_momentum_exchange(
                    alpha[i], rho_g, rho_l,
                    v_g_axial[i], v_l_axial[i],
                    v_g_adj, v_l_adj,
                    conn_params['gap_length'],
                    conn_params['gap_width']
                )
                
                # 计算湍流交混
                w_t = self.physics.calculate_turbulent_mixing(
                    self.physics.BETA,
                    conn_params['gap_width'],
                    rho_g, rho_l,
                    v_g_axial[i], v_l_axial[i]
                )
                
                M_t = self.physics.calculate_turbulent_mixing_momentum(
                    w_t, v_g_axial[i], v_g_adj
                )
                
                M_lat_g_total += M_lat_g + M_t
                M_lat_l_total += M_lat_l + M_t
            
            # 更新轴向速度
            v_g_axial[i] += dt * (
                -dp_dz/rho_g + F_w_g + F_i/alpha[i] + 
                M_lat_g_total/(rho_g * geom['area']) - self.physics.g
            )
            
            v_l_axial[i] += dt * (
                -dp_dz/rho_l + F_w_l - F_i/(1-alpha[i]) + 
                M_lat_l_total/(rho_l * geom['area']) - self.physics.g
            )
        
        # 更新横向速度（使用横向压力梯度）
        for adj_ch in adjacent_channels:
            conn_params = subchannel.get_connection_parameters(adj_ch)
            dp_dx = (adj_ch.pressure - pressure) / conn_params['centroid_distance']
            v_g_lateral += dt * (-dp_dx/rho_g)
            v_l_lateral += dt * (-dp_dx/rho_l)
        
        # 更新空泡份额（考虑漂移速度的影响）
        j_drift = alpha * (1 - alpha) * v_gj
        alpha_new = alpha + dt * (
            -np.gradient(alpha * v_g_axial, self.grid.dz) +
            np.gradient(j_drift, self.grid.dz)
        )
        
        # 限制空泡份额在物理范围内
        alpha_new = np.clip(alpha_new, 0.0, 1.0)
        
        # 更新结果
        subchannel.alpha = alpha_new
        subchannel.velocity['vapor']['axial'] = v_g_axial
        subchannel.velocity['liquid']['axial'] = v_l_axial
        subchannel.velocity['vapor']['lateral'] = v_g_lateral
        subchannel.velocity['liquid']['lateral'] = v_l_lateral
        
        # 计算残差
        residuals = {
            'velocity_g': np.max(np.abs(v_g_axial - subchannel.previous_velocity['vapor']['axial'])),
            'velocity_l': np.max(np.abs(v_l_axial - subchannel.previous_velocity['liquid']['axial'])),
            'alpha': np.max(np.abs(alpha_new - subchannel.previous_alpha))
        }
        
        return residuals
    
    def solve_time_step(self):
        """使用PKM方法求解一个时间步"""
        # 计算时间步长
        dt = self.calculate_time_step()
        
        # 使用PKM求解器
        converged = self.pkm_solver.solve_step(self.grid.subchannels, dt)
        
        if not converged:
            # 减小时间步长重试
            self.dt /= 2
            if self.dt < self.min_dt:
                raise RuntimeError("计算发散：时间步长小于最小允许值")
            return False
            
        return True

    def solve_improved_piso(self, subchannel):
        """使用改进的PISO算法求解速度和压力场
    
        Args:
            subchannel: 子通道对象
        """
        # 第一步：速度预测
        self.predict_velocity_piso(subchannel)
        
        # 第二步：压力修正（使用Rhie-Chow插值）
        self.solve_pressure_correction_rhie_chow(subchannel)
        
        # 第三步：考虑界面动量交换的压力修正
        self.solve_pressure_correction_with_interfacial(subchannel)
        
        # 第四步：速度修正
        self.correct_velocity_with_interfacial(subchannel)

    def solve_pressure_correction_rhie_chow(self, subchannel):
        """使用Rhie-Chow插值求解压力修正方程"""
        # 获取物性参数
        n = len(subchannel.pressure)
        
        # 构建系数矩阵
        A = np.zeros((n, n))
        b = np.zeros(n)
        
        # 边界条件
        A[0, 0] = 1.0
        b[0] = 0.0  # 入口压力修正为0
        A[-1, -1] = 1.0
        b[-1] = 0.0  # 出口压力修正为0
        
        # 内部节点
        for i in range(1, n-1):
            try:
                # 获取物性（添加错误检查）
                props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
                props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
                
                # 确保密度为正值
                rho_g = max(props_g['density'], 1e-6)
                rho_l = max(props_l['density'], 1e-6)
                
                # 确保空泡份额在物理范围内
                alpha = np.clip(subchannel.alpha[i], 1e-6, 1-1e-6)
                
                # 计算系数
                ap_g = rho_g / self.dt
                ap_l = rho_l / self.dt
                
                # 添加数值稳定项
                stability_factor = 1e-8
                
                # 构建方程系数
                A[i, i-1] = -(alpha/rho_g + (1-alpha)/rho_l) / (self.grid.dz**2)
                A[i, i] = 2*(alpha/rho_g + (1-alpha)/rho_l) / (self.grid.dz**2) + stability_factor
                A[i, i+1] = -(alpha/rho_g + (1-alpha)/rho_l) / (self.grid.dz**2)
                
                # 计算源项（质量不平衡）
                mass_imbalance = self.calculate_mass_imbalance(subchannel, i)
                if np.isfinite(mass_imbalance):
                    b[i] = mass_imbalance
                else:
                    self.log_and_print(f"警告：节点 {i} 的质量不平衡计算结果无效", level='warning')
                    b[i] = 0.0
                    
            except Exception as e:
                self.log_and_print(f"警告：节点 {i} 的压力修正方程系数计算失败: {str(e)}", level='warning')
                # 使用简单的拉普拉斯算子作为后备
                A[i, i-1] = -1.0
                A[i, i] = 2.0
                A[i, i+1] = -1.0
                b[i] = 0.0
        
        try:
            # 增加对角优势以提高稳定性
            for i in range(n):
                A[i,i] = max(A[i,i], 1.1 * np.sum(np.abs(A[i,:])))
            
            # 求解压力修正方程
            dp = np.linalg.solve(A, b)
            
            # 检查解的有效性
            if not np.all(np.isfinite(dp)):
                raise ValueError("压力修正解包含无效值")
            
            # 限制修正幅度
            max_correction = 0.1 * np.mean(np.abs(subchannel.pressure))
            dp = np.clip(dp, -max_correction, max_correction)
            
            # 应用压力修正（使用较小的松弛因子）
            omega = 0.2  # 更小的松弛因子
            subchannel.pressure += omega * dp
            
            # 确保压力保持在物理范围内
            min_pressure = 1e5  # 最小压力（1 bar）
            subchannel.pressure = np.maximum(subchannel.pressure, min_pressure)
            
        except Exception as e:
            self.log_and_print(f"警告：压力修正方程求解失败: {str(e)}", level='warning')
            # 使用简单的压力梯度估计
            dp = (subchannel.pressure[-1] - subchannel.pressure[0]) / (n - 1)
            subchannel.pressure = np.linspace(subchannel.pressure[0], subchannel.pressure[-1], n)

    def calculate_mass_imbalance(self, subchannel, i):
        """计算质量不平衡
        
        Args:
            subchannel: 子通道对象
            i: 当前节点索引
            
        Returns:
            float: 质量不平衡
        """
        try:
            # 获取物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 确保密度为正值
            rho_g = max(props_g['density'], 1e-6)
            rho_l = max(props_l['density'], 1e-6)
            
            # 确保空泡份额在物理范围内
            alpha = np.clip(subchannel.alpha[i], 1e-6, 1-1e-6)
            alpha_prev = np.clip(subchannel.alpha[i-1], 1e-6, 1-1e-6)
            
            # 计算速度项（添加限制以避免过大的速度）
            v_g = np.clip(subchannel.velocity['vapor']['axial'][i], -10.0, 10.0)
            v_l = np.clip(subchannel.velocity['liquid']['axial'][i], -10.0, 10.0)
            v_g_prev = np.clip(subchannel.velocity['vapor']['axial'][i-1], -10.0, 10.0)
            v_l_prev = np.clip(subchannel.velocity['liquid']['axial'][i-1], -10.0, 10.0)
            
            # 计算当前质量流量
            G_g = alpha * rho_g * v_g
            G_l = (1 - alpha) * rho_l * v_l
            
            # 计算上一个节点的质量流量
            G_g_prev = alpha_prev * rho_g * v_g_prev
            G_l_prev = (1 - alpha_prev) * rho_l * v_l_prev
            
            # 计算质量流量的变化率
            dG_g_dz = (G_g - G_g_prev) / self.grid.dz
            dG_l_dz = (G_l - G_l_prev) / self.grid.dz
            
            # 计算总的质量不平衡
            mass_imbalance = -(dG_g_dz + dG_l_dz)
            
            # 添加数值稳定项
            stability_factor = 1e-8
            mass_imbalance += stability_factor * (subchannel.pressure[i] - subchannel.previous_pressure[i])
            
            # 限制质量不平衡的大小
            max_imbalance = 1.0  # 最大允许的质量不平衡
            mass_imbalance = np.clip(mass_imbalance, -max_imbalance, max_imbalance)
            
            return mass_imbalance if np.isfinite(mass_imbalance) else 0.0
            
        except Exception as e:
            self.log_and_print(f"警告：节点 {i} 的质量不平衡计算失败: {str(e)}", level='warning')
            return 0.0

    def calculate_velocity_residual(self, subchannel):
        """计算速度残差
        
        Args:
            subchannel: 子通道对象
            
        Returns:
            float: 速度残差
        """
        n_axial = len(subchannel.alpha)
        velocity_residual = 0.0
        
        for i in range(1, n_axial-1):
            velocity_residual += np.abs(subchannel.velocity['vapor']['axial'][i] - 
                                      subchannel.previous_velocity['vapor']['axial'][i])
        
        return velocity_residual

    def calculate_pressure_residual(self, subchannel):
        """计算压力残差
        
        Args:
            subchannel: 子通道对象
            
        Returns:
            float: 压力残差
        """
        n_axial = len(subchannel.alpha)
        pressure_residual = 0.0
        
        for i in range(1, n_axial-1):
            pressure_residual += np.abs(subchannel.pressure[i] - 
                                      subchannel.previous_pressure[i])
        
        return pressure_residual

    def predict_velocity_piso(self, subchannel):
        """PISO算法的速度预测步骤"""
        n_axial = len(subchannel.alpha)
        
        for i in range(1, n_axial-1):
            # 获取物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 计算对流项（使用迎风格式）
            for phase, props in [('vapor', props_g), ('liquid', props_l)]:
                v = subchannel.velocity[phase]['axial']
                rho = props['density']
                
                # 计算对流通量
                if v[i] > 0:
                    conv_flux = v[i] * (v[i] - v[i-1]) / self.grid.dz
                else:
                    conv_flux = v[i] * (v[i+1] - v[i]) / self.grid.dz
                
                # 计算压力梯度
                dp_dz = (subchannel.pressure[i+1] - subchannel.pressure[i-1]) / (2 * self.grid.dz)
                
                # 预测速度（考虑对流项）
                subchannel.velocity[phase]['axial'][i] = (
                    subchannel.previous_velocity[phase]['axial'][i] -
                    self.dt * (conv_flux + dp_dz / rho)
                )

    def solve_pressure_correction_with_interfacial(self, subchannel):
        """考虑相间动量交换的压力修正方程求解"""
        n_axial = len(subchannel.alpha)
        
        try:
            # 构建压力修正方程系数矩阵
            A = np.zeros((n_axial, n_axial))
            b = np.zeros(n_axial)
            
            # 边界条件
            A[0, 0] = 1.0
            b[0] = 0.0  # 入口压力修正为0
            A[-1, -1] = 1.0
            b[-1] = 0.0  # 出口压力修正为0
            
            # 内部节点
            for i in range(1, n_axial-1):
                try:
                    # 获取物性
                    props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
                    props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
                    
                    # 确保密度为正值
                    rho_g = max(props_g['density'], 1e-6)
                    rho_l = max(props_l['density'], 1e-6)
                    
                    # 确保空泡份额在物理范围内
                    alpha = np.clip(subchannel.alpha[i], 1e-6, 1-1e-6)
                    
                    # 计算速度（添加限制）
                    v_g = np.clip(subchannel.velocity['vapor']['axial'][i], -10.0, 10.0)
                    v_l = np.clip(subchannel.velocity['liquid']['axial'][i], -10.0, 10.0)
                    v_r = abs(v_g - v_l)  # 相对速度
                    
                    # 计算相间动量交换系数（改进模型）
                    C_i = min(0.44 * alpha * (1 - alpha) * rho_l * v_r / self.grid.dz, 1e3)
                    
                    # 计算系数（考虑数值稳定性）
                    a_g = alpha / rho_g
                    a_l = (1 - alpha) / rho_l
                    
                    # 添加数值稳定项
                    stability_factor = 1e-6
                    
                    # 修改系数矩阵（考虑相间动量交换和数值稳定性）
                    A[i, i-1] = -(a_g + a_l) / (self.grid.dz ** 2)
                    A[i, i] = 2 * (a_g + a_l) / (self.grid.dz ** 2) + C_i + stability_factor
                    A[i, i+1] = -(a_g + a_l) / (self.grid.dz ** 2)
                    
                    # 计算源项（包含相间动量交换）
                    mass_imbalance = self.calculate_mass_imbalance(subchannel, i)
                    b[i] = mass_imbalance + C_i * (v_g - v_l)
                    
                except Exception as e:
                    self.log_and_print(f"警告：节点 {i} 的压力修正方程系数计算失败: {str(e)}", level='warning')
                    # 使用简单的拉普拉斯算子作为后备
                    A[i, i-1] = -1.0
                    A[i, i] = 2.0 + stability_factor
                    A[i, i+1] = -1.0
                    b[i] = 0.0
            
            # 增加对角优势以提高稳定性
            for i in range(n_axial):
                A[i,i] = max(A[i,i], 1.1 * np.sum(np.abs(A[i,:])))
            
            # 求解压力修正方程
            pressure_correction = np.linalg.solve(A, b)
            
            # 检查解的有效性
            if not np.all(np.isfinite(pressure_correction)):
                raise ValueError("压力修正解包含无效值")
            
            # 限制修正幅度
            max_correction = 0.1 * np.mean(np.abs(subchannel.pressure))
            pressure_correction = np.clip(pressure_correction, -max_correction, max_correction)
            
            # 应用压力修正（使用较小的松弛因子）
            omega = 0.1  # 更小的松弛因子
            subchannel.pressure += omega * pressure_correction
            
            # 确保压力保持在物理范围内
            min_pressure = 1e5  # 最小压力（1 bar）
            subchannel.pressure = np.maximum(subchannel.pressure, min_pressure)
            
        except Exception as e:
            self.log_and_print(f"警告：压力修正方程求解失败: {str(e)}", level='warning')
            # 使用简单的压力梯度估计
            dp = (subchannel.pressure[-1] - subchannel.pressure[0]) / (n_axial - 1)
            subchannel.pressure = np.linspace(subchannel.pressure[0], subchannel.pressure[-1], n_axial)

    def correct_velocity_with_interfacial(self, subchannel):
        """考虑相间动量交换的速度修正"""
        n_axial = len(subchannel.alpha)
        
        for i in range(1, n_axial-1):
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 计算压力梯度
            dp_dz = (subchannel.pressure[i+1] - subchannel.pressure[i-1]) / (2 * self.grid.dz)
            
            # 计算相间动量交换
            v_g = subchannel.velocity['vapor']['axial'][i]
            v_l = subchannel.velocity['liquid']['axial'][i]
            v_r = abs(v_g - v_l)
            
            alpha = subchannel.alpha[i]
            rho_l = props_l['density']
            
            # 相间动量交换系数
            C_i = 0.44 * alpha * (1 - alpha) * rho_l * v_r / self.grid.dz
            
            # 速度修正（考虑相间动量交换）
            for phase, props in [('vapor', props_g), ('liquid', props_l)]:
                velocity_correction = (
                    -dp_dz / props['density'] +
                    C_i * (v_g - v_l) / props['density']
                )
                
                subchannel.velocity[phase]['axial'][i] += (
                    self.relaxation_factors['velocity'] * velocity_correction
                ) 

    def update_void_fraction_tvd(self, subchannel):
        """使用TVD格式更新空泡份额
        
        Args:
            subchannel: 子通道对象
        """
        n_axial = len(subchannel.alpha)
        alpha_new = np.zeros(n_axial)
        
        # 获取物性
        props_g = self.physics.calculate_properties(subchannel.pressure[0], PhaseType.VAPOR)
        props_l = self.physics.calculate_properties(subchannel.pressure[0], PhaseType.LIQUID)
        
        # 计算漂移速度
        v_gj, C0 = self.physics.calculate_drift_velocity(props_g['density'], 
                                                       props_l['density'])
        
        # 初始化漂移通量数组
        j_drift_array = np.zeros(n_axial)
        
        # 内部节点计算
        for i in range(1, n_axial-1):
            # 1. 计算TVD限制器
            r_plus = self.calculate_tvd_ratio(subchannel.alpha, i, 'forward')
            r_minus = self.calculate_tvd_ratio(subchannel.alpha, i, 'backward')
            
            # 使用van Leer限制器
            phi_plus = self.van_leer_limiter(r_plus)
            phi_minus = self.van_leer_limiter(r_minus)
            
            # 2. 计算通量
            # 2.1 计算混合速度
            j = (subchannel.alpha[i] * subchannel.velocity['vapor']['axial'][i] + 
                 (1-subchannel.alpha[i]) * subchannel.velocity['liquid']['axial'][i])
            
            # 2.2 计算实际气相速度（考虑漂移速度）
            v_g = C0 * j + v_gj
            
            # 2.3 计算漂移通量
            j_drift = subchannel.alpha[i] * (1-subchannel.alpha[i]) * v_gj
            j_drift_array[i] = j_drift
            
            # 3. 计算通量限制
            if v_g > 0:
                flux_plus = v_g * (subchannel.alpha[i] + 
                                  0.5 * phi_plus * (subchannel.alpha[i+1] - subchannel.alpha[i]))
                flux_minus = v_g * (subchannel.alpha[i-1] + 
                                   0.5 * phi_minus * (subchannel.alpha[i] - subchannel.alpha[i-1]))
            else:
                flux_plus = v_g * (subchannel.alpha[i+1] - 
                                  0.5 * phi_plus * (subchannel.alpha[i+1] - subchannel.alpha[i]))
                flux_minus = v_g * (subchannel.alpha[i] - 
                                   0.5 * phi_minus * (subchannel.alpha[i] - subchannel.alpha[i-1]))
            
            # 4. 计算漂移通量的导数
            if subchannel.previous_drift_flux is None:
                d_drift_dz = 0.0
            else:
                d_drift_dz = (j_drift - subchannel.previous_drift_flux[i]) / self.dt
            
            # 5. 更新空泡份额
            alpha_new[i] = subchannel.alpha[i] - self.dt/self.grid.dz * (flux_plus - flux_minus) + self.dt * d_drift_dz
        
        # 边界条件处理
        alpha_new[0] = subchannel.alpha[0]  # 入口边界条件
        alpha_new[-1] = alpha_new[-2]       # 出口外推
        
        # 物理限制
        alpha_new = np.clip(alpha_new, 0.0, 1.0)
            
        # 更新空泡份额
        subchannel.alpha = alpha_new
        
        # 保存当前漂移通量用于下一步计算
        subchannel.previous_drift_flux = j_drift_array

    def calculate_tvd_ratio(self, alpha: np.ndarray, i: int, direction: str) -> float:
        """计算TVD比率
        
        Args:
            alpha: 空泡份额数组
            i: 当前节点索引
            direction: 计算方向 ('forward' 或 'backward')
            
        Returns:
            float: TVD比率
        """
        eps = 1e-10  # 防止除零
        
        if direction == 'forward':
            if i < len(alpha)-2:
                numerator = alpha[i+2] - alpha[i+1]
                denominator = alpha[i+1] - alpha[i]
            else:
                numerator = 0.0
                denominator = 1.0
        else:  # backward
            if i > 1:
                numerator = alpha[i] - alpha[i-1]
                denominator = alpha[i-1] - alpha[i-2]
            else:
                numerator = 0.0
                denominator = 1.0
        
        return numerator / (denominator + eps)

    def van_leer_limiter(self, r: float) -> float:
        """van Leer限制器
        
        Args:
            r: TVD比率
            
        Returns:
            float: 限制器值
        """
        return (r + abs(r)) / (1 + abs(r))

    def minmod_limiter(self, r: float) -> float:
        """minmod限制器
        
        Args:
            r: TVD比率
            
        Returns:
            float: 限制器值
        """
        return max(0.0, min(1.0, r))

    def superbee_limiter(self, r: float) -> float:
        """superbee限制器
        
        Args:
            r: TVD比率
            
        Returns:
            float: 限制器值
        """
        return max(0.0, min(2.0*r, 1.0), min(r, 2.0)) 

    def solve_lateral_flow(self, subchannel, adjacent_channels, dt: float):
        """求解横向流动
        
        Args:
            subchannel: 当前子通道
            adjacent_channels: 相邻子通道列表
            dt: 时间步长
        """
        n_axial = len(subchannel.alpha)
        
        # 获取几何参数
        geom = subchannel.get_geometry_parameters()
        
        for i in range(1, n_axial-1):
            # 获取当前节点的物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 当前通道的参数
            alpha = subchannel.alpha[i]
            rho_g = props_g['density']
            rho_l = props_l['density']
            v_g = subchannel.velocity['vapor']['axial'][i]
            v_l = subchannel.velocity['liquid']['axial'][i]
            
            # 遍历相邻通道
            for adj_ch in adjacent_channels:
                # 获取连接参数
                conn_params = subchannel.get_connection_parameters(adj_ch)
                gap_width = conn_params['gap_width']
                gap_length = conn_params['gap_length']
                
                # 计算横向压力梯度
                dp_dy = (adj_ch.pressure[i] - subchannel.pressure[i]) / conn_params['centroid_distance']
                
                # 计算横向流动阻力系数（KIJ = 0.5）
                K_ij = 0.5
                
                # 计算横向流动速度（考虑FTM = 0.8）
                FTM = 0.8  # 横向动量因子
                v_g_lat = subchannel.velocity['vapor']['lateral'][i]
                v_l_lat = subchannel.velocity['liquid']['lateral'][i]
                
                # 计算横向流动阻力
                R_g = -K_ij * rho_g * v_g_lat * abs(v_g_lat) / gap_width
                R_l = -K_ij * rho_l * v_l_lat * abs(v_l_lat) / gap_width
                
                # 计算湍流混合项（β = 0.065）
                beta = 0.065  # 湍流交混系数
                w_t = self.physics.calculate_turbulent_mixing(
                    beta, gap_width, rho_g, rho_l, v_g, v_l
                )
                
                # 计算横向动量交换
                M_lat_g, M_lat_l = self.physics.calculate_lateral_momentum_exchange(
                    alpha, rho_g, rho_l, v_g, v_l,
                    adj_ch.velocity['vapor']['axial'][i],
                    adj_ch.velocity['liquid']['axial'][i],
                    gap_length, gap_width
                )
                
                # 计算空隙漂移效应
                v_gj, C0 = self.physics.calculate_drift_velocity(rho_g, rho_l)
                j_drift_lat = alpha * (1 - alpha) * v_gj * FTM
                
                # 更新横向速度
                # 气相
                dv_g_lat = dt * (
                    -dp_dy/rho_g +  # 压力梯度项
                    R_g/rho_g +     # 横向阻力项
                    M_lat_g/(rho_g * alpha * gap_width) +  # 横向动量交换项
                    w_t/(rho_g * alpha * gap_width)  # 湍流混合项
                )
                
                # 液相
                dv_l_lat = dt * (
                    -dp_dy/rho_l +  # 压力梯度项
                    R_l/rho_l +     # 横向阻力项
                    M_lat_l/(rho_l * (1-alpha) * gap_width) +  # 横向动量交换项
                    w_t/(rho_l * (1-alpha) * gap_width)  # 湍流混合项
                )
                
                # 应用松弛因子并更新横向速度
                subchannel.velocity['vapor']['lateral'][i] += (
                    self.relaxation_factors['velocity'] * dv_g_lat
                )
                subchannel.velocity['liquid']['lateral'][i] += (
                    self.relaxation_factors['velocity'] * dv_l_lat
                )
                
                # 更新质量通量（考虑横向流动的影响）
                G_g_lat = alpha * rho_g * subchannel.velocity['vapor']['lateral'][i]
                G_l_lat = (1-alpha) * rho_l * subchannel.velocity['liquid']['lateral'][i]
                
                # 保存当前时间步的横向质量通量
                subchannel.previous_mass_flux['vapor'][i] = G_g_lat
                subchannel.previous_mass_flux['liquid'][i] = G_l_lat

    def predict_velocity_piso(self, subchannel):
        """PISO算法的速度预测步骤"""
        n_axial = len(subchannel.alpha)
        
        for i in range(1, n_axial-1):
            # 获取物性
            props_g = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(subchannel.pressure[i], PhaseType.LIQUID)
            
            # 计算对流项（使用迎风格式）
            for phase, props in [('vapor', props_g), ('liquid', props_l)]:
                v = subchannel.velocity[phase]['axial']
                rho = props['density']
                
                # 计算对流通量
                if v[i] > 0:
                    conv_flux = v[i] * (v[i] - v[i-1]) / self.grid.dz
                else:
                    conv_flux = v[i] * (v[i+1] - v[i]) / self.grid.dz
                
                # 计算压力梯度
                dp_dz = (subchannel.pressure[i+1] - subchannel.pressure[i-1]) / (2 * self.grid.dz)
                
                # 预测速度（考虑对流项）
                subchannel.velocity[phase]['axial'][i] = (
                    subchannel.previous_velocity[phase]['axial'][i] -
                    self.dt * (conv_flux + dp_dz / rho)
                ) 

class PKMSolver:
    """Picard-Krylov求解器"""
    def __init__(self, grid, physics_model):
        self.grid = grid
        self.physics = physics_model
        
        # Picard迭代参数
        self.max_picard_iterations = 50  # 增加最大迭代次数
        self.picard_tolerance = 0.05  # 放宽收敛判据
        self.backup_tolerance = 0.1   # 放宽备用收敛判据
        
        # GMRES参数
        self.gmres_tolerance = 1e-4   # 放宽GMRES收敛判据
        self.max_gmres_iterations = 200
        self.gmres_restart = 30       # 减小重启参数
        
        # 松弛因子（更保守）
        self.relaxation_factors = {
            'pressure': 0.1,          # 减小压力松弛因子
            'velocity': 0.3,          # 减小速度松弛因子
            'void_fraction': 0.1      # 减小空泡份额松弛因子
        }
        
        # 数值限制参数
        self.max_velocity = 50.0      # m/s
        self.min_pressure = 1e5       # Pa
        self.max_pressure = 1e7       # Pa
        
        # 初始化预处理器
        self.preconditioner = PhysicsBasedPreconditioner()
        
    def solve_step(self, subchannels, dt):
        """执行单个时间步的PKM求解
        
        Args:
            subchannels: 子通道列表
            dt: 时间步长
            
        Returns:
            bool: 是否收敛
        """
        for iter in range(self.max_picard_iterations):
            max_residual = 0.0
            
            # 1. 求解动量方程获得质量流量
            momentum_residual = self._solve_momentum_equations(subchannels, dt)
            max_residual = max(max_residual, momentum_residual)
            
            # 2. 使用PGMRES求解压力修正方程
            pressure_residual = self._solve_pressure_correction(subchannels)
            max_residual = max(max_residual, pressure_residual)
            
            # 3. 更新两相本构方程
            self._update_constitutive_relations(subchannels)
            
            # 4. 应用数值稳定化
            self._apply_numerical_stabilization(subchannels)
            
            # 5. 收敛性检查
            if self._check_convergence(iter, max_residual):
                return True
                
        # 如果不收敛，返回False（外部处理时间步长调整）
        return False
        
    def _solve_momentum_equations(self, subchannels, dt):
        """求解动量方程（添加安全检查）"""
        max_residual = 0.0
        
        for subchannel in subchannels:
            n_cells = len(subchannel.pressure)
            
            for i in range(1, n_cells-1):
                # 计算压力梯度（添加限制）
                dp = subchannel.pressure[i+1] - subchannel.pressure[i-1]
                dp = np.clip(dp, -1e5, 1e5)  # 限制压力差
                dp_dz = dp / (2 * self.grid.dz)
                
                # 获取物性（添加压力限制）
                p = np.clip(subchannel.pressure[i], self.min_pressure, self.max_pressure)
                props_g = self.physics.calculate_properties(p, PhaseType.VAPOR)
                props_l = self.physics.calculate_properties(p, PhaseType.LIQUID)
                
                # 限制速度
                for phase in ['vapor', 'liquid']:
                    v = subchannel.velocity[phase]['axial'][i]
                    subchannel.velocity[phase]['axial'][i] = np.clip(v, -self.max_velocity, self.max_velocity)
                
                # 计算摩擦力（添加安全检查）
                try:
                    F_wall_g = self.physics.calculate_wall_friction(
                        PhaseType.VAPOR,
                        props_g['density'],
                        subchannel.velocity['vapor']['axial'][i],
                        subchannel.hydraulic_diameter,
                        props_g['viscosity']
                    )
                    F_wall_l = self.physics.calculate_wall_friction(
                        PhaseType.LIQUID,
                        props_l['density'],
                        subchannel.velocity['liquid']['axial'][i],
                        subchannel.hydraulic_diameter,
                        props_l['viscosity']
                    )
                    F_interface = self.physics.calculate_interfacial_friction(
                        subchannel.alpha[i],
                        props_g['density'],
                        props_l['density'],
                        subchannel.velocity['vapor']['axial'][i],
                        subchannel.velocity['liquid']['axial'][i],
                        subchannel.hydraulic_diameter
                    )
                except (ValueError, RuntimeWarning) as e:
                    print(f"警告：摩擦力计算出错 - {str(e)}")
                    F_wall_g = F_wall_l = F_interface = 0.0
                
                # 限制摩擦力
                max_friction = 1e4  # 最大摩擦力
                F_wall_g = np.clip(F_wall_g, -max_friction, max_friction)
                F_wall_l = np.clip(F_wall_l, -max_friction, max_friction)
                F_interface = np.clip(F_interface, -max_friction, max_friction)
                
                # 更新速度场（带松弛因子和限制）
                for phase, props, F_wall in [
                    ('vapor', props_g, F_wall_g),
                    ('liquid', props_l, F_wall_l)
                ]:
                    # 计算新速度
                    new_velocity = (
                        subchannel.velocity[phase]['axial'][i] -
                        dt * (
                            dp_dz / max(props['density'], 1.0) +  # 防止除零
                            F_wall +
                            (F_interface if phase == 'vapor' else -F_interface)
                        )
                    )
                    
                    # 限制速度变化
                    max_change = 0.5 * abs(subchannel.velocity[phase]['axial'][i])
                    delta_v = new_velocity - subchannel.velocity[phase]['axial'][i]
                    delta_v = np.clip(delta_v, -max_change, max_change)
                    new_velocity = subchannel.velocity[phase]['axial'][i] + delta_v
                    
                    # 应用松弛因子和速度限制
                    subchannel.velocity[phase]['axial'][i] = np.clip(
                        self.relaxation_factors['velocity'] * new_velocity +
                        (1 - self.relaxation_factors['velocity']) * 
                        subchannel.velocity[phase]['axial'][i],
                        -self.max_velocity,
                        self.max_velocity
                    )
                    
                    # 更新残差（添加安全检查）
                    if not np.isnan(new_velocity):
                        residual = abs(delta_v)
                        max_residual = max(max_residual, residual)
                    
        return max_residual
        
    def _solve_pressure_correction(self, subchannels):
        """改进的压力修正求解"""
        max_residual = 0.0
        
        for subchannel in subchannels:
            # 构建压力修正方程组
            A, b = self._build_pressure_correction_system(subchannel)
            
            # 应用预处理器
            M = self.preconditioner.build_matrix(A, subchannel)
            
            # 确保矩阵是稀疏的
            from scipy.sparse import csr_matrix
            A_sparse = csr_matrix(A)
            M_sparse = csr_matrix(M)
            
            # 使用GMRES求解
            from scipy.sparse.linalg import gmres, LinearOperator
            
            # 创建预处理器操作符
            def M_operator(x):
                return np.linalg.solve(M_sparse.toarray(), x)
            
            M_op = LinearOperator(M_sparse.shape, M_operator)
            
            try:
                delta_p, info = gmres(
                    A_sparse, b,
                    M=M_op,
                    atol=self.gmres_tolerance,
                    maxiter=self.max_gmres_iterations,
                    restart=self.gmres_restart
                )
                
                if info > 0:
                    print(f"警告：GMRES在{info}次迭代后未收敛")
                    # 使用直接求解作为备选
                    try:
                        delta_p = np.linalg.solve(A, b)
                    except np.linalg.LinAlgError:
                        print("警告：直接求解失败，使用近似解")
                        delta_p = np.zeros_like(b)
                elif info < 0:
                    print(f"警告：GMRES出现错误，错误码：{info}")
                    delta_p = np.zeros_like(b)
                
                # 限制压力修正的大小
                max_correction = 0.1 * np.mean(subchannel.pressure)
                delta_p = np.clip(delta_p, -max_correction, max_correction)
                
                # 更新压力场（带松弛因子和限制）
                for i in range(len(subchannel.pressure)):
                    new_pressure = subchannel.pressure[i] + delta_p[i]
                    subchannel.pressure[i] = np.clip(
                        self.relaxation_factors['pressure'] * new_pressure +
                        (1 - self.relaxation_factors['pressure']) * subchannel.pressure[i],
                        self.min_pressure,
                        self.max_pressure
                    )
                    
                    # 更新残差
                    if not np.isnan(delta_p[i]):
                        residual = abs(delta_p[i] / max(subchannel.pressure[i], 1.0))
                        max_residual = max(max_residual, residual)
                    
            except Exception as e:
                print(f"GMRES求解出错：{str(e)}")
                raise
                
        return max_residual
        
    def _build_pressure_correction_system(self, subchannel):
        """构建压力修正方程组"""
        n_cells = len(subchannel.pressure)
        A = np.zeros((n_cells, n_cells))
        b = np.zeros(n_cells)
        
        # 设置边界条件
        A[0, 0] = 1.0
        A[-1, -1] = 1.0
        
        # 构建内部节点方程
        for i in range(1, n_cells-1):
            # 获取物性
            props_g = self.physics.calculate_properties(
                subchannel.pressure[i], PhaseType.VAPOR)
            props_l = self.physics.calculate_properties(
                subchannel.pressure[i], PhaseType.LIQUID)
            
            # 计算系数
            alpha = np.clip(subchannel.alpha[i], 1e-6, 1-1e-6)
            coeff = (alpha / props_g['density'] + (1 - alpha) / props_l['density'])
            
            # 设置矩阵元素
            A[i, i-1] = -coeff / (self.grid.dz ** 2)
            A[i, i] = 2 * coeff / (self.grid.dz ** 2)
            A[i, i+1] = -coeff / (self.grid.dz ** 2)
            
            # 计算源项（质量不平衡）
            G_g = alpha * props_g['density'] * subchannel.velocity['vapor']['axial'][i]
            G_l = (1 - alpha) * props_l['density'] * subchannel.velocity['liquid']['axial'][i]
            
            dG_dz = (G_g + G_l - subchannel.previous_mass_flux['vapor'][i] - 
                    subchannel.previous_mass_flux['liquid'][i]) / self.grid.dz
            
            b[i] = -dG_dz
            
        return A, b
        
    def _update_constitutive_relations(self, subchannels):
        """更新两相本构方程"""
        for subchannel in subchannels:
            n_cells = len(subchannel.pressure)
            
            for i in range(n_cells):
                # 更新质量通量
                props_g = self.physics.calculate_properties(
                    subchannel.pressure[i], PhaseType.VAPOR)
                props_l = self.physics.calculate_properties(
                    subchannel.pressure[i], PhaseType.LIQUID)
                
                alpha = subchannel.alpha[i]
                subchannel.previous_mass_flux['vapor'][i] = (
                    alpha * props_g['density'] * 
                    subchannel.velocity['vapor']['axial'][i]
                )
                subchannel.previous_mass_flux['liquid'][i] = (
                    (1 - alpha) * props_l['density'] * 
                    subchannel.velocity['liquid']['axial'][i]
                )
                
    def _apply_numerical_stabilization(self, subchannels):
        """应用数值稳定化技术"""
        for subchannel in subchannels:
            # 限制空泡份额
            subchannel.alpha = np.clip(subchannel.alpha, 1e-6, 1-1e-6)
            
            # 确保压力为正
            subchannel.pressure = np.maximum(subchannel.pressure, 1e5)
            
            # 限制速度变化
            for phase in ['vapor', 'liquid']:
                for direction in ['axial', 'lateral']:
                    v = subchannel.velocity[phase][direction]
                    v_max = np.max(np.abs(v)) * 1.5
                    subchannel.velocity[phase][direction] = np.clip(v, -v_max, v_max)
                    
    def _check_convergence(self, iter_num, max_residual):
        """检查收敛性"""
        # 主收敛判据
        if max_residual < self.picard_tolerance:
            return True
            
        # 备用收敛判据
        if (iter_num >= 20 and max_residual < self.backup_tolerance):
            return True
            
        return False

class PhysicsBasedPreconditioner:
    """基于物理的预处理器"""
    def __init__(self):
        self.matrix = None
        
    def build_matrix(self, system_matrix, subchannel):
        """构建预处理器矩阵，仅保留中心单元和轴向相邻单元的元素"""
        n = system_matrix.shape[0]
        self.matrix = np.zeros_like(system_matrix)
        
        for i in range(n):
            # 保留主对角线元素
            self.matrix[i,i] = system_matrix[i,i]
            
            # 保留轴向相邻元素
            if i > 0:
                self.matrix[i,i-1] = system_matrix[i,i-1]
            if i < n-1:
                self.matrix[i,i+1] = system_matrix[i,i+1]
                
        return self.matrix 